/*
 * SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES.
 * All rights reserved. SPDX-License-Identifier: LicenseRef-NvidiaProprietary
 *
 * NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
 * property and proprietary rights in and to this material, related
 * documentation and any modifications thereto. Any use, reproduction,
 * disclosure or distribution of this material and related documentation
 * without an express license agreement from NVIDIA CORPORATION or
 * its affiliates is strictly prohibited.
 */

#include "msda_plugin.h"

#if NV_TENSORRT_MAJOR < 10

#include <cstring>

#include "cuda_fp16.h"

using namespace nvinfer1;
using namespace plugin;

namespace nvinfer1 {
namespace plugin {
namespace {
void throwPluginError(char const* file, char const* func, int line, int,
                      char const* cond) {
  std::stringstream ss;
  ss << file << ":" << line << ": " << cond << " check failed";
  throw std::runtime_error(ss.str().c_str());
};
void reportAssertion(char const* assertion, char const* file, int line) {
  std::cerr << file << ":" << line << ": assertion " << assertion << " failed";
  std::abort();
};
}  // namespace
#define PLUGIN_VALIDATE(condition)                                   \
  {                                                                  \
    if (!(condition)) {                                              \
      throwPluginError(__FILE__, __func__, __LINE__, 0, #condition); \
    }                                                                \
  }

#define PLUGIN_ASSERT(assertion)                       \
  {                                                    \
    if (!(assertion)) {                                \
      reportAssertion(#assertion, __FILE__, __LINE__); \
    }                                                  \
  }

MSDAPlugin::MSDAPlugin(const int2* spatial_shape, const float* scale_value,
                       float scale_offset, float scale_weight, float scale_out,
                       Layout source_layout) {
#pragma unroll
  for (int s = 0; s < MSDA_MAX_SCALES; s++) {
    params.spatial_shape[s] = spatial_shape[s];
    params.scale_value[s] = scale_value[s];
  }
  params.scale_offset = scale_offset;
  params.scale_weight = scale_weight;
  params.scale_out = scale_out;
  params.source_layout = source_layout;
}

MSDAPlugin::MSDAPlugin(void const* data, size_t length) {
  this->params = *reinterpret_cast<const Params*>(data);
}

nvinfer1::IPluginV2DynamicExt* MSDAPlugin::clone() const PLUGIN_NOEXCEPT {
  try {
    MSDAPlugin* plugin = new MSDAPlugin(
        params.spatial_shape, params.scale_value, params.scale_offset,
        params.scale_weight, params.scale_out, params.source_layout);
    plugin->setPluginNamespace(getPluginNamespace());
    return plugin;
  } catch (const std::exception& e) {
    std::cout << e.what() << std::endl;
    // caughtError(e);
  }
  return nullptr;
}

nvinfer1::DimsExprs MSDAPlugin::getOutputDimensions(
    int32_t outputIndex, nvinfer1::DimsExprs const* inputs, int32_t nbInputs,
    nvinfer1::IExprBuilder& exprBuilder) PLUGIN_NOEXCEPT {
  nvinfer1::DimsExprs ret;
  ret.nbDims = 4;
  ret.d[0] = inputs[0].d[0];
  ret.d[1] = inputs[0].d[1];
  ret.d[2] = inputs[0].d[2];
  switch (params.source_layout) {
    case Layout::CONCAT_NHWC:
      ret.d[3] = exprBuilder.operation(nvinfer1::DimensionOperation::kFLOOR_DIV,
                                       *inputs[2].d[2], *inputs[0].d[2]);
      break;
    case Layout::CONCAT_NCHW32:
      ret.d[3] = exprBuilder.operation(
          nvinfer1::DimensionOperation::kFLOOR_DIV,
          *exprBuilder.operation(nvinfer1::DimensionOperation::kPROD,
                                 *inputs[2].d[1], *inputs[2].d[3]),
          *inputs[0].d[2]);
      break;
    default:  // NCHW32
              // NHWC
      ret.d[3] = exprBuilder.operation(nvinfer1::DimensionOperation::kFLOOR_DIV,
                                       *inputs[2].d[1], *inputs[0].d[2]);
      break;
  }
  return ret;
}

bool MSDAPlugin::supportsFormatCombination(
    int32_t pos, nvinfer1::PluginTensorDesc const* inOut, int32_t nbInputs,
    int32_t nbOutputs) PLUGIN_NOEXCEPT {
  int nbValues = nbInputs - 2;
  PLUGIN_ASSERT((nbValues < MSDA_MAX_SCALES));
  PLUGIN_ASSERT((nbOutputs == 1));

  if (pos == 0)
    return (inOut[pos].type == nvinfer1::DataType::kINT8 ||
            inOut[pos].type == nvinfer1::DataType::kHALF) &&
           inOut[pos].format == nvinfer1::TensorFormat::kLINEAR;
  if (pos == 2) {
    switch (params.source_layout) {
      case Layout::NHWC:
        return (inOut[pos].type == inOut[0].type) &
               (inOut[pos].format == nvinfer1::TensorFormat::kHWC8 ||
                inOut[pos].format == nvinfer1::TensorFormat::kHWC16);
      case Layout::NCHW32:
        return (inOut[pos].type == inOut[0].type) &
               (inOut[pos].format == nvinfer1::TensorFormat::kCHW32);
      default:  // CONCAT_NHWC:
                // CONCAT_NCHW32:
        return (inOut[pos].type == inOut[0].type) &
               (inOut[pos].format == nvinfer1::TensorFormat::kLINEAR);
    }
  }
  if (pos > 2 && pos < nbInputs) {
    return inOut[pos].type == inOut[0].type &&
           inOut[pos].format == inOut[2].format;
  }
  return inOut[pos].type == inOut[0].type &&
         inOut[pos].format == nvinfer1::TensorFormat::kLINEAR;
}

void MSDAPlugin::configurePlugin(
    nvinfer1::DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
    nvinfer1::DynamicPluginTensorDesc const* outputs,
    int32_t nbOutputs) PLUGIN_NOEXCEPT {
  int nbValues = nbInputs - 2;
  const int batch_size = inputs[0].desc.dims.d[0];
  const int num_queries = inputs[0].desc.dims.d[1];
  const int num_heads = inputs[0].desc.dims.d[2];
  const int num_scales = inputs[0].desc.dims.d[3];
  const int num_points_per_scale = inputs[0].desc.dims.d[4];
  const int num_channels_per_head = outputs[0].desc.dims.d[3];

  PLUGIN_ASSERT(num_points_per_scale % 4 == 0);
  // todo: relax this
  if (params.source_layout == Layout::NHWC)
    PLUGIN_ASSERT((num_heads * num_channels_per_head) % 16 == 0);

  // Check for valid input dimensions
  PLUGIN_ASSERT(inputs[0].desc.dims.nbDims == 6);
  PLUGIN_ASSERT(inputs[1].desc.dims.nbDims == 5);

  // Check batch_size consistency
  PLUGIN_ASSERT(inputs[0].desc.dims.d[0] == batch_size);
  PLUGIN_ASSERT(inputs[1].desc.dims.d[0] == batch_size);

  // Check num_queries consistency
  PLUGIN_ASSERT(inputs[0].desc.dims.d[1] == num_queries);
  PLUGIN_ASSERT(inputs[1].desc.dims.d[1] == num_queries);

  // Check num_heads consistency
  PLUGIN_ASSERT(inputs[0].desc.dims.d[2] == num_heads);
  PLUGIN_ASSERT(inputs[1].desc.dims.d[2] == num_heads);

  // Check num_scales consistency
  PLUGIN_ASSERT(inputs[0].desc.dims.d[3] == num_scales);
  PLUGIN_ASSERT(inputs[1].desc.dims.d[3] == num_scales);

  // Check num_points_per_scale consistency
  PLUGIN_ASSERT(inputs[0].desc.dims.d[4] == num_points_per_scale);
  PLUGIN_ASSERT(inputs[1].desc.dims.d[4] == num_points_per_scale);

  // sampling locations' last dim must be == 2 (xy)
  PLUGIN_ASSERT(inputs[0].desc.dims.d[5] == 2);

  int total_spatial_locations = 0;
  if (params.source_layout == Layout::CONCAT_NHWC ||
      params.source_layout == Layout::CONCAT_NCHW32) {
    for (int s = 0; s < num_scales; s++) {
      auto s_s = params.spatial_shape[s];
      PLUGIN_ASSERT(s_s.x > 0 && s_s.y > 0);
      total_spatial_locations += s_s.x * s_s.y;
    }
  }

  switch (params.source_layout) {
    case Layout::CONCAT_NHWC:
      PLUGIN_ASSERT(nbValues == 1);
      PLUGIN_ASSERT(inputs[2].desc.dims.nbDims == 3);
      PLUGIN_ASSERT(inputs[2].desc.dims.d[0] == batch_size);
      PLUGIN_ASSERT(inputs[2].desc.dims.d[1] == total_spatial_locations);
      PLUGIN_ASSERT(inputs[2].desc.dims.d[2] ==
                    num_channels_per_head * num_heads);
      break;
    case Layout::CONCAT_NCHW32:
      PLUGIN_ASSERT(nbValues == 1);
      PLUGIN_ASSERT(inputs[2].desc.dims.nbDims == 4);
      PLUGIN_ASSERT(inputs[2].desc.dims.d[0] == batch_size);
      PLUGIN_ASSERT(inputs[2].desc.dims.d[1] * 32 ==
                    num_channels_per_head * num_heads);
      PLUGIN_ASSERT(inputs[2].desc.dims.d[2] == total_spatial_locations);
      PLUGIN_ASSERT(inputs[2].desc.dims.d[3] == 32);
      break;
    default:  // NHWC
              // NCHW32
      PLUGIN_ASSERT(nbValues == num_scales);
      for (int s = 0; s < nbValues; s++) {
        PLUGIN_ASSERT(inputs[s + 2].desc.dims.nbDims == 4);
        PLUGIN_ASSERT(inputs[s + 2].desc.dims.d[0] == batch_size);
        PLUGIN_ASSERT(inputs[s + 2].desc.dims.d[1] ==
                      num_channels_per_head * num_heads);
      }
      break;
  }
}

size_t MSDAPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs,
                                    int32_t nbInputs,
                                    nvinfer1::PluginTensorDesc const* outputs,
                                    int32_t nbOutputs) const PLUGIN_NOEXCEPT {
  return 0;
}

int32_t MSDAPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc,
                            nvinfer1::PluginTensorDesc const* outputDesc,
                            void const* const* inputs, void* const* outputs,
                            void* workSpace,
                            cudaStream_t stream) PLUGIN_NOEXCEPT {
  const void* data_multiscale_value[MSDA_MAX_SCALES];
  void* data_output = outputs[0];
  const void* data_sampling_offsets = inputs[0];
  const void* data_attn_weight = inputs[1];

  const int batch_size = inputDesc[0].dims.d[0];
  const int num_queries = inputDesc[0].dims.d[1];
  const int num_heads = inputDesc[0].dims.d[2];
  const int num_scales = inputDesc[0].dims.d[3];
  const int num_points_per_scale = inputDesc[0].dims.d[4];
  const int num_channels_per_head = outputDesc[0].dims.d[3];

  params.scale_offset =
      params.scale_offset > 0 ? params.scale_offset : inputDesc[0].scale;
  params.scale_weight =
      params.scale_weight > 0 ? params.scale_weight : inputDesc[1].scale;
  params.scale_out =
      params.scale_out > 0 ? params.scale_out : outputDesc[0].scale;

  switch (params.source_layout) {
    case Layout::NHWC:
    case Layout::NCHW32:
      for (int s = 0; s < num_scales; s++) {
        params.spatial_shape[s] =
            make_int2(inputDesc[s + 2].dims.d[2], inputDesc[s + 2].dims.d[3]);
        data_multiscale_value[s] = inputs[s + 2];
        params.scale_value[s] = params.scale_value[s] > 0
                                    ? params.scale_value[s]
                                    : inputDesc[s + 2].scale;
      }
      break;
    default:  //  CONCAT_NHWC
              // CONCAT_NCHW32
      data_multiscale_value[0] = inputs[2];
      params.scale_value[0] = params.scale_value[0] > 0 ? params.scale_value[0]
                                                        : inputDesc[2].scale;
      break;
  }

  decltype(&msda<int8_t, Layout::NCHW32, 16>) func = nullptr;
  if (inputDesc[0].type == nvinfer1::DataType::kINT8) {
    switch (params.source_layout) {
      case Layout::NCHW32: {
        if (num_channels_per_head % 16 == 0) {
          func = &msda<int8_t, Layout::NCHW32, 16>;
        } else if (num_channels_per_head % 8 == 0) {
          func = &msda<int8_t, Layout::NCHW32, 8>;
        } else {
          func = &msda<int8_t, Layout::NCHW32, 4>;
        }
        break;
      }
      default: {
        // NHWC
        // CONCAT_NCHW32
        // CONCAT_NHWC
        return -1;
      }
    }
  } else if (inputDesc[0].type == nvinfer1::DataType::kHALF) {
    switch (params.source_layout) {
      case Layout::NHWC: {
        if (num_channels_per_head % 8 == 0) {
          func = &msda<half, Layout::NHWC, 8>;
        } else if (num_channels_per_head % 4 == 0) {
          func = &msda<half, Layout::NHWC, 4>;
        }
        break;
      }
      case Layout::NCHW32: {
        if (num_channels_per_head % 8 == 0) {
          func = &msda<half, Layout::NCHW32, 8>;
        } else if (num_channels_per_head % 4 == 0) {
          func = &msda<half, Layout::NCHW32, 4>;
        }
        break;
      }
      case Layout::CONCAT_NCHW32: {
        if (num_channels_per_head % 8 == 0) {
          func = &msda<half, Layout::CONCAT_NCHW32, 8>;
        } else if (num_channels_per_head % 4 == 0) {
          func = &msda<half, Layout::CONCAT_NCHW32, 4>;
        }
        break;
      }
      case Layout::CONCAT_NHWC: {
        if (num_channels_per_head % 8 == 0) {
          func = &msda<half, Layout::CONCAT_NHWC, 8>;
        } else if (num_channels_per_head % 4 == 0) {
          func = &msda<half, Layout::CONCAT_NHWC, 4>;
        }
        break;
      }
      default: {
        return -1;
      }
    }
  }

  if (!func) return -1;

  return func(stream, data_output, data_multiscale_value, data_sampling_offsets,
              data_attn_weight, params.spatial_shape, batch_size, num_queries,
              num_heads, num_scales, num_points_per_scale,
              num_channels_per_head, params.scale_value, params.scale_offset,
              params.scale_weight, params.scale_out);
}

void MSDAPlugin::attachToContext(
    cudnnContext* cudnnContext, cublasContext* cublasContext,
    nvinfer1::IGpuAllocator* gpuAllocator) PLUGIN_NOEXCEPT {}

void MSDAPlugin::detachFromContext() PLUGIN_NOEXCEPT {}

// IPluginV2Ext Methods
nvinfer1::DataType MSDAPlugin::getOutputDataType(
    int32_t index, nvinfer1::DataType const* inputTypes,
    int32_t nbInputs) const PLUGIN_NOEXCEPT {
  return inputTypes[0];
}

// IPluginV2 Methods
char const* MSDAPlugin::getPluginType() const PLUGIN_NOEXCEPT {
  return MSDA_NAME;
}

char const* MSDAPlugin::getPluginVersion() const PLUGIN_NOEXCEPT {
  return MSDA_VERSION;
}

int32_t MSDAPlugin::getNbOutputs() const PLUGIN_NOEXCEPT { return 1; }

int32_t MSDAPlugin::initialize() PLUGIN_NOEXCEPT { return 0; }

void MSDAPlugin::terminate() PLUGIN_NOEXCEPT {}

size_t MSDAPlugin::getSerializationSize() const PLUGIN_NOEXCEPT {
  return sizeof(Params);
}

void MSDAPlugin::serialize(void* buffer) const PLUGIN_NOEXCEPT {
  *reinterpret_cast<Params*>(buffer) = params;
}

void MSDAPlugin::destroy() PLUGIN_NOEXCEPT { delete this; }

void MSDAPlugin::setPluginNamespace(char const* pluginNamespace)
    PLUGIN_NOEXCEPT {
  mNamespace = pluginNamespace;
}
char const* MSDAPlugin::getPluginNamespace() const PLUGIN_NOEXCEPT {
  return mNamespace.c_str();
}

// Pluginv1 Creator

MSDAPluginCreator::MSDAPluginCreator() {
  mPluginAttributes.clear();

  mPluginAttributes.emplace_back(nvinfer1::PluginField("scale_value"));
  mPluginAttributes.emplace_back(nvinfer1::PluginField("scale_offset"));
  mPluginAttributes.emplace_back(nvinfer1::PluginField("scale_weight"));
  mPluginAttributes.emplace_back(nvinfer1::PluginField("scale_out"));
  mPluginAttributes.emplace_back(nvinfer1::PluginField("feature_layout"));
  mPluginAttributes.emplace_back(nvinfer1::PluginField("spatial_shape"));

  mFC.nbFields = mPluginAttributes.size();
  mFC.fields = mPluginAttributes.data();
}

char const* MSDAPluginCreator::getPluginName() const PLUGIN_NOEXCEPT {
  return MSDA_NAME;
}

char const* MSDAPluginCreator::getPluginVersion() const PLUGIN_NOEXCEPT {
  return MSDA_VERSION;
}

nvinfer1::PluginFieldCollection const* MSDAPluginCreator::getFieldNames()
    PLUGIN_NOEXCEPT {
  return &mFC;
}

IPluginV2* MSDAPluginCreator::createPlugin(
    char const* name, PluginFieldCollection const* fc) PLUGIN_NOEXCEPT {
  try {
    float scale_value[MSDA_MAX_SCALES], scale_offset = -1.0f,
                                        scale_weight = -1.0f, scale_out = -1.0f;
    int2 spatial_shape[MSDA_MAX_SCALES];
    for (auto& s_v : scale_value) s_v = -1.0f;
    for (auto& s_s : spatial_shape) s_s = {0, 0};
    Layout source_layout = Layout::NCHW32;
    for (int i = 0; i < fc->nbFields; i++) {
      auto name = fc->fields[i].name;
      auto type = fc->fields[i].type;
      auto length = fc->fields[i].length;
      if (!strcmp(name, "feature_layout")) {
        PLUGIN_VALIDATE(type == PluginFieldType::kCHAR);
        auto feature_layout = static_cast<char const*>(fc->fields[i].data);
        if (!strcmp(feature_layout, "NCHW32")) {
          source_layout = Layout::NCHW32;
        } else if (!strcmp(feature_layout, "NCHW")) {
          source_layout = Layout::NCHW32;
        } else if (!strcmp(feature_layout, "NHWC")) {
          source_layout = Layout::NHWC;
        } else if (!strcmp(feature_layout, "CONCAT_NCHW32")) {
          source_layout = Layout::CONCAT_NCHW32;
        } else if (!strcmp(feature_layout, "CONCAT_NHWC")) {
          source_layout = Layout::CONCAT_NHWC;
        } else {
          throwPluginError(__FILE__, __func__, __LINE__, 0,
                           "does not support this format!");
        }
      }
      if (!strcmp(name, "spatial_shape")) {
        PLUGIN_VALIDATE(type == PluginFieldType::kINT32);
        for (int s = 0; s * 2 < length; s++) {
          spatial_shape[s].x =
              reinterpret_cast<int const*>(fc->fields[i].data)[2 * s];
          spatial_shape[s].y =
              reinterpret_cast<int const*>(fc->fields[i].data)[2 * s + 1];
        }
      }
      if (!strcmp(name, "scale_value")) {
        PLUGIN_VALIDATE(type == PluginFieldType::kFLOAT32 ||
                        type == PluginFieldType::kFLOAT64);
        for (int s = 0; s < length; s++) {
          scale_value[s] =
              type == PluginFieldType::kFLOAT32
                  ? reinterpret_cast<float const*>(fc->fields[i].data)[s]
                  : reinterpret_cast<double const*>(fc->fields[i].data)[s];
        }
      }
      if (!strcmp(name, "scale_offset")) {
        PLUGIN_VALIDATE(type == PluginFieldType::kFLOAT32 ||
                        type == PluginFieldType::kFLOAT64);
        scale_offset =
            type == PluginFieldType::kFLOAT32
                ? *(reinterpret_cast<float const*>(fc->fields[i].data))
                : *(reinterpret_cast<double const*>(fc->fields[i].data));
      }
      if (!strcmp(name, "scale_weight")) {
        PLUGIN_VALIDATE(type == PluginFieldType::kFLOAT32 ||
                        type == PluginFieldType::kFLOAT64);
        scale_weight =
            type == PluginFieldType::kFLOAT32
                ? *(reinterpret_cast<float const*>(fc->fields[i].data))
                : *(reinterpret_cast<double const*>(fc->fields[i].data));
      }
      if (!strcmp(name, "scale_out")) {
        PLUGIN_VALIDATE(type == PluginFieldType::kFLOAT32 ||
                        type == PluginFieldType::kFLOAT64);
        scale_out =
            type == PluginFieldType::kFLOAT32
                ? *(reinterpret_cast<float const*>(fc->fields[i].data))
                : *(reinterpret_cast<double const*>(fc->fields[i].data));
      }
    }
    MSDAPlugin* plugin =
        new MSDAPlugin(spatial_shape, scale_value, scale_offset, scale_weight,
                       scale_out, source_layout);
    return plugin;
  } catch (const std::exception& e) {
    printf(e.what());
  }
  return nullptr;
}

IPluginV2* MSDAPluginCreator::deserializePlugin(
    char const* name, void const* serialData,
    size_t serialLength) PLUGIN_NOEXCEPT {
  try {
    auto plugin = new MSDAPlugin(serialData, serialLength);
    plugin->setPluginNamespace(getPluginNamespace());

    return plugin;
  } catch (const std::exception& e) {
    printf("%s\n", e.what());
  }
  return nullptr;
}

void MSDAPluginCreator::setPluginNamespace(char const* pluginNamespace)
    PLUGIN_NOEXCEPT {
  mNamespace = pluginNamespace;
}

char const* MSDAPluginCreator::getPluginNamespace() const PLUGIN_NOEXCEPT {
  return mNamespace.c_str();
}

REGISTER_TENSORRT_PLUGIN(MSDAPluginCreator);

}  // namespace plugin
}  // namespace nvinfer1
#endif
